package com.zyx.flinkdemo.stream.watermark;

import com.zyx.flinkdemo.pojo.WordTs;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;

/**
 * @author Yaxi.Zhang
 * @since 2021/5/27 15:34
 * desc:
 */
public class WatermarkDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<String> source = env.socketTextStream("localhost", 7777);
        //设置水印设置周期
        env.getConfig().setAutoWatermarkInterval(200);
        env.setParallelism(1);
        SingleOutputStreamOperator<WordTs> map = source.map(new MapFunction<String, WordTs>() {
            @Override
            public WordTs map(String value) throws Exception {
                String[] split = value.split(",");
                return new WordTs(split[0], Long.parseLong(split[1]));
            }
        });
        // 生成水印 event-允许最大延迟的时间=watermark
       /* SingleOutputStreamOperator<WordTs> wordTsDs = map.assignTimestampsAndWatermarks(MyWatermarkStrategy
                .<WordTs>forMyWatermarkGenerator(Duration.ofSeconds(3L))
                .withTimestampAssigner(new SerializableTimestampAssigner<WordTs>() {
                    @Override
                    public long extractTimestamp(WordTs element, long recordTimestamp) {
                        return element.getTs();
                    }
                }));*/

        SingleOutputStreamOperator<WordTs> wordTsDs = map.assignTimestampsAndWatermarks(WatermarkStrategy
                .<WordTs>forBoundedOutOfOrderness(Duration.ofSeconds(1L))
                .withTimestampAssigner(new SerializableTimestampAssigner<WordTs>() {
                    @Override
                    public long extractTimestamp(WordTs element, long recordTimestamp) {
                        return element.getTs();
                    }
                }));

        //wordTsDs.print();


        // 处理数据
        wordTsDs
                .keyBy(new KeySelector<WordTs, String>() {
                    @Override
                    public String getKey(WordTs value) throws Exception {
                        return value.getWord();
                    }
                })
                .process(new KeyedProcessFunction<String, WordTs, String>() {

                    private ValueState<Long> tumblingTimerState;
                    private ValueState<Integer> count;

                    @Override
                    public void open(Configuration parameters) throws Exception {
                        tumblingTimerState = getRuntimeContext()
                                .getState(new ValueStateDescriptor<>("tumbling_timer_state", Long.class));
                        count = getRuntimeContext()
                                .getState(new ValueStateDescriptor<Integer>("count_value_dtate", Integer.class));
                    }

                    @Override
                    public void processElement(WordTs value,
                                               Context ctx,
                                               Collector<String> out) throws Exception {
                        long curWatermark = ctx.timerService().currentWatermark();

                        if (tumblingTimerState.value() == null) {
                            long onTimerTimestamp = getOnTimerTimestamp(curWatermark);
                            ctx.timerService().registerEventTimeTimer(onTimerTimestamp);
                            tumblingTimerState.update(onTimerTimestamp);
                        }
                        if (count.value() == null) {
                            count.update(1);
                        } else {
                            count.update(count.value() + 1);
                        }
                    }

                    @Override
                    public void onTimer(long timestamp, OnTimerContext ctx, Collector<String> out) throws Exception {
                        StringBuilder sb = new StringBuilder();
                        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
                        Instant instant = Instant.ofEpochMilli(timestamp);
                        String ts = formatter.format(instant.atZone(ZoneId.systemDefault()));

                        sb
                                .append("ts >>> ").append(ts).append(" word >>>").append(ctx.getCurrentKey())
                                .append(", count >>> ").append(count.value());
                        out.collect(sb.toString());
                        tumblingTimerState.clear();
                        count.clear();
                    }

                    private long getOnTimerTimestamp(long curWatermark) {
                        // watermark = maxTs - 1ms
                        long maxTs = curWatermark + 1L;
                        return maxTs - (maxTs % 5000) + 5000;
                    }
                })
                .print();
        env.execute();
    }
}
